import numpy as np
from math import sin, cos, radians, sqrt, degrees
import matplotlib.pyplot as plt

# --- I. 物理常数和弹丸参数 ---
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 绝热垂直温度梯度 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 干燥空气比气体常数 (J/kg·K)
P_TERM = (G / (R * L)) - 1 # 用于计算密度比的指数
C_SOUND = 338.0      # 海平面声速 (m/s)

M = 5.0              # 弹丸质量 (kg)
A = 0.0095           # 弹丸截面积 (m^2)
MAGNUS_CL = 0.0      # 马格努斯升力系数：已设为 0.0，取消马格努斯效应
# 注意: 虽然 omega_x 仍存在，但由于 MAGNUS_CL=0，马格努斯加速度为零。

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) # 马赫数
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) # 阻力系数

# --- II. 动态物理模型和RK4函数 ---
def get_drag_coefficient(v_rel):
    """根据相对速度计算空气阻力系数 (Cd)。"""
    mach_number = v_rel / C_SOUND
    return np.interp(mach_number, MA_POINTS, CD_POINTS)

def get_air_density(z):
    """根据高度 Z 计算空气密度。"""
    h = max(0.0, z) 
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0: return 0.0 
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

def derivatives_of_state(S, v_wind_y, omega_x, params):
    """计算状态向量 S 在时间上的微分 (加速度)。"""
    x, y, z, vx, vy, vz = S
    M, A, G = params 
    
    # 计算相对速度 (相对空气)
    vx_rel = vx
    vy_rel = vy - v_wind_y
    vz_rel = vz
    v_rel = sqrt(vx_rel**2 + vy_rel**2 + vz_rel**2)

    if v_rel < 1e-6:
        # 弹丸静止或速度极低时
        return np.array([vx, vy, vz, 0.0, 0.0, -G])

    rho = get_air_density(z)
    Cd = get_drag_coefficient(v_rel)

    # 1. 空气阻力加速度
    drag_coeff_factor = -(Cd * rho * A * v_rel) / (2.0 * M)
    dvx_drag = drag_coeff_factor * vx_rel
    dvy_drag = drag_coeff_factor * vy_rel
    dvz_drag = drag_coeff_factor * vz_rel

    # 2. 马格努斯力加速度 (由于 MAGNUS_CL=0.0，此加速度为零)
    W = np.array([omega_x, 0.0, 0.0])         
    V_rel = np.array([vx_rel, vy_rel, vz_rel]) 
    magnus_cross_product = np.cross(W, V_rel)

    magnus_accel_factor = (0.5 * MAGNUS_CL * rho * A * v_rel) / M
    a_magnus = magnus_accel_factor * (magnus_cross_product / v_rel) # a_magnus 约等于 [0, 0, 0]

    # 合并加速度
    dvx_dt = dvx_drag + a_magnus[0]
    dvy_dt = dvy_drag + a_magnus[1]
    dvz_dt = dvz_drag - G + a_magnus[2]

    return np.array([vx, vy, vz, dvx_dt, dvy_dt, dvz_dt])

def rk4_step(S, dt, vwy, omega, params):
    """四阶龙格-库塔方法进行一步积分。"""
    f = lambda S: derivatives_of_state(S, vwy, omega, params)
    
    K1 = dt * f(S)
    K2 = dt * f(S + K1 / 2.0)
    K3 = dt * f(S + K2 / 2.0)
    K4 = dt * f(S + K3)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

def simulate_and_get_impact(v0, theta_deg, psi_deg, omega_x, dt, v_wind_y):
    """模拟 3D 轨迹并返回最终落点 (R, Y) 坐标。"""
    params = (M, A, G)
    theta_rad = radians(theta_deg)
    psi_rad = radians(psi_deg)
    
    v_horiz = v0 * cos(theta_rad)
    S = np.array([
        0.0, 0.0, 0.0, 
        v_horiz * cos(psi_rad), 
        v_horiz * sin(psi_rad), 
        v0 * sin(theta_rad)
    ])
    
    while S[2] >= -0.01:
        S_new = rk4_step(S, dt, v_wind_y, omega_x, params) 
        
        if S_new[2] < 0 and S[2] >= 0:
            # 线性插值计算冲击点
            t_impact = S[2] / (S[2] - S_new[2])
            R_final = S[0] + (S_new[0] - S[0]) * t_impact
            Y_final = S[1] + (S_new[1] - S[1]) * t_impact
            return R_final, Y_final
            
        S = S_new
        
    return S[0], S[1]

# --- III. 搜索函数 ---
def binary_search_v0(target_range, theta_deg, omega_x, dt, V_low=10.0, V_high=2000.0, tolerance=1.0):
    """查找固定角度和目标距离所需的 V0 (无风假设)。"""
    
    for _ in range(50):
        V_mid = (V_low + V_high) / 2.0
        # 假设 0 初始偏航 (psi=0), 0 横风 (v_wind_y=0)
        R_mid, _ = simulate_and_get_impact(V_mid, theta_deg, 0.0, omega_x, dt, 0.0)
        
        if abs(R_mid - target_range) < tolerance:
            return V_mid
        
        if R_mid < target_range:
            V_low = V_mid
        else:
            V_high = V_mid
            
    raise ValueError(f"V0 search failed to converge for target {target_range}m.")

def binary_search_yaw_correction(v0, theta_deg, target_range, v_wind_y, omega_x, dt, yaw_low=-10.0, yaw_high=10.0, tolerance=0.1):
    """查找所需的初始偏航角 (psi_deg) 以使 Y_final 接近 0。"""
    
    for _ in range(50):
        psi_mid = (yaw_low + yaw_high) / 2.0
        R_final, Y_final = simulate_and_get_impact(v0, theta_deg, psi_mid, omega_x, dt, v_wind_y)

        if abs(Y_final) < tolerance:
            return psi_mid
        
        # 如果 Y_final > 0，说明落点在目标右侧，需要更小的偏航角 (即向左瞄准)
        if Y_final > 0:
            yaw_high = psi_mid
        # 如果 Y_final < 0，说明落点在目标左侧，需要更大的偏航角 (即向右瞄准)
        else:
            yaw_low = psi_mid
            
    return np.nan

# --- IV. 主运行和表格输出 ---

def generate_windage_table():
    
    # --- 默认/固定参数 ---
    THETA = 35.0          # 固定发射角
    OMEGA_X = 500.0       # 自旋角速度 (虽然 Magnus 被禁用，但保留变量)
    DT = 0.01             # RK4 步长 (s)

    # --- 表格参数 (直线距离 1000m-1500m) ---
    R_MIN = 1000.0
    R_MAX = 1500.0
    R_STEP = 100.0        # 直线距离步长
    V_W_MIN = -30.0       # 最小横风速度 (向左吹)
    V_W_MAX = 30.0        # 最大横风速度 (向右吹)
    V_W_STEP = 5.0        # 横风速度步长
    
    R_TARGETS = np.arange(R_MIN, R_MAX + R_STEP, R_STEP)
    V_W_TARGETS = np.arange(V_W_MIN, V_W_MAX + V_W_STEP, V_W_STEP)

    if len(R_TARGETS) < 2 or len(V_W_TARGETS) < 2:
        print("❌ 错误: 步长设置不合理，无法生成表格。")
        return

    # 存储 V0 查找表
    V0_LOOKUP = {}
    
    # 存储结果的二维数组
    results_table = np.zeros((len(V_W_TARGETS), len(R_TARGETS)))

    print("--- 🎯 正在进行 V₀ 和校正角双重计算 (已禁用马格努斯效应) ---")
    
    # 1. 外层循环：遍历目标距离 R_target (列)
    for j, R_target in enumerate(R_TARGETS):
        print(f"\n[R={R_target:.0f}m] -> 正在搜索所需 V₀...")
        
        try:
            # 为每个 R_target 独立搜索最佳 V0_R
            V0_R = binary_search_v0(R_target, THETA, OMEGA_X, DT)
            V0_LOOKUP[R_target] = V0_R
            print(f"  ✅ V₀ 找到: {V0_R:.3f} m/s. 正在计算横风校正...")
            
        except ValueError:
            print(f"  ❌ V₀ 搜索失败，跳过该距离.")
            results_table[:, j] = np.nan
            continue

        # 2. 内层循环：遍历横风速度 V_w (行)
        for i, V_w in enumerate(V_W_TARGETS):
            
            # 使用 V0_R (特定距离的V0) 来计算校正角
            psi_correction = binary_search_yaw_correction(
                v0=V0_R, 
                theta_deg=THETA, 
                target_range=R_target,
                v_wind_y=V_w, 
                omega_x=OMEGA_X, 
                dt=DT
            )
            results_table[i, j] = psi_correction

    # --- 结果表格输出 ---
    print("\n" + "="*90)
    print("                      🎯 弹道横风校正角查表 (单位: 度) 🎯")
    print("                   [注意: 已禁用马格努斯效应，仅考虑阻力与横风]")
    print(f"  **固定发射角**: θ={THETA}°, RK4 dt={DT}s")
    print("="*90)

    # 格式化表头 (X轴：目标直线距离)
    header = " | " + " | ".join([f"{r:^5.0f}" for r in R_TARGETS]) + " |"
    print(" 横风 V_y (m/s) \\ 目标距离 R (m)")
    print("-" * len(header))
    print(header)
    print("-" * len(header))
    
    # 格式化数据行 (Y轴：横风速度)
    for i, V_w in enumerate(V_W_TARGETS):
        row = [f"{V_w:^10.1f}"]
        for j in range(len(R_TARGETS)):
            val = results_table[i, j]
            if np.isnan(val):
                row.append(" N/A ")
            else:
                row.append(f"{val:>+5.2f}") # 输出带正负号，两位小数
        
        print(f" | {V_w:>+10.1f} | " + " | ".join(row[1:]) + " |")

    print("="*90)
    print("V₀ Lookup (m/s):", ", ".join([f"{r:.0f}m: {v:.1f}" for r, v in V0_LOOKUP.items()]))
    print("💡 查表指引: 正角度 (+) 表示需要向目标右侧瞄准 (抵消左吹风)。")
    
if __name__ == "__main__":
    generate_windage_table()